#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Feb 16 21:16:42 2022

@author: liqingsimac
"""

#5.2.3. 
'''
import numpy as np
import matplotlib.pyplot as plt

# use the next command in IPython terminal mode
#plt.ion()

# use the next command in IPython non-inline notebook mode
#%matplotlib

# use the next command in IPython inline notebook mode
#%matplotlib notebook

x = np.linspace(-np.pi, np.pi, 101)
y = np.sin(x) + np.sin(3*x)/3.0

#plt.plot(x,y)
#plt.xlabel('x')
#plt.ylabel('y')
#plt.title('A simple plot')
#plt.show()
#plt.savefig('pic/fig-5-2-3.png')

fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(x,y)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_title('A simple plot')
#fig.savefig('pic/fig-5-2-3.png')
'''

#5.3.
'''
>> %matplotlib?
>> %matplotlib
>> %matplotlib inline
'''

#5.4.5.
'''
import numpy as np
import matplotlib.pyplot as plt

x=np.linspace(-np.pi,np.pi,101)
f=np.ones_like(x)
f[x<0]=-1
y1=(4/np.pi)*(np.sin(x)+np.sin(3*x)/3.0)
y2=y1+(4/np.pi)*(np.sin(5*x)/5.0+np.sin(7*x)/7.0)
y3=y2+(4/np.pi)*(np.sin(9*x)/9.0+np.sin(11*x)/11.0)

fig=plt.figure()
ax=fig.add_subplot(111)
ax.plot(x,f,'b-',lw=3,label='f(x)')
ax.plot(x,y1,'c--',lw=2,label='two terms')
ax.plot(x,y2,'r-.',lw=2,label='four terms')
ax.plot(x,y3,'b:',lw=2,label='six terms')
ax.legend(loc='best')
ax.set_xlabel('x',style='italic')
ax.set_ylabel('partial sums',style='italic')
fig.subtitle('Partial sums for Fourier series of f(x)',
             size=16,weight='bold')
fig.savefig('pic/fig-5-4-5.png')
'''

#5.5.
'''
import numpy as np
import matplotlib.pyplot as plt

theta=np.linspace(0,2*np.pi,201)
r1=np.abs(np.cos(5.0*theta)-1.5*np.sin(3.0*theta))
r2=theta/np.pi
r3=2.25*np.ones_like(theta)

fig=plt.figure()
ax=fig.add_subplot(111,projection='polar')
ax.plot(theta,r1,label='trig')
ax.plot(5*theta,r2,label='spiral')
ax.plot(theta,r3,label='circle')
ax.legend(loc='best')
'''

#5.6. error bar
'''
import numpy as np
import numpy.random as npr

x=np.linspace(0,4,21)
y=np.exp(-x)
xe=0.08*npr.randn(len(x))
ye=0.1*npr.randn(len(y))

import matplotlib.pyplot as plt
fig=plt.figure()
ax=fig.add_subplot(111)
ax.errorbar(x,y,fmt='bo',lw=2,xerr=xe,yerr=ye,
            ecolor='r',elinewidth=1)
fig.savefig('pic/fig-5-6.png')
'''

#5.7.
'''
import numpy as np
x=np.linspace(0,2,101)
y=(x-1)**3+1
import matplotlib.pyplot as plt

fig=plt.figure()
ax=fig.add_subplot(111)
ax.plot(x,y)
ax.annotate('point of inplection at x=1',xy=(1,1),
            xytext=(0.8,0.5),
            arrowprops=dict(facecolor='black',
                            width=1,
                            shrink=0.05))
'''

#5.9.
'''
import numpy as np
import matplotlib.pyplot as plt
fig=plt.figure()
ax=fig.add_subplot(111)
[X,Y]=np.mgrid[-3:3:61j,-3:3:61j]
Z=X**2-Y**2
#curves=ax.contour(X,Y,Z,12,colors='b')
#ax.clabel(curves)
im=ax.contourf(X,Y,Z,12)
fig.colorbar(im,orientation='vertical')
fig.suptitle(r'The level contour of $z=x^2-y^2$',fontsize=20)
'''

#5.10.
'''
import numpy as np
import matplotlib.pyplot as plt
x=np.linspace(0,5,101)
y1=1.0/(x+1.0)
y2=np.exp(-x)
y3=np.exp(-0.1*x**2)
y4=np.exp(-5*x**2)

fig=plt.figure()
ax1=fig.add_subplot(221)
ax1.plot(x,y1)
ax2=fig.add_subplot(222)
ax2.plot(x,y2)
ax3=fig.add_subplot(223)
ax3.plot(x,y3)
ax4=fig.add_subplot(224)
ax4.plot(x,y4)
fig.suptitle('Various decay functions')
'''

#5.11. Mandelbrot Set
#'''
import numpy as np
from time import time

#Set the parameters
max_iter=256  #maximum number of iterations
nx,ny=1024,1024   #x- and y-image resolution
#x_lo,x_hi=-2.0,1.0  #x bounds in complex plain
#y_lo,y_hi=-1.5,1.5  #y bounds in complex plain
x_lo,x_hi=-2.0,-1.9998  #x bounds in complex plain
y_lo,y_hi=-0.0001,0.0001  #y bounds in complex plain
start_time=time()

#Construct the two-dimensional arrays
ix,iy=np.mgrid[0:nx,0:ny]
x,y=np.mgrid[x_lo:x_hi:1j*nx, y_lo:y_hi:1j*ny]
c=x+1j*y
#holds pixel rgb data
esc_parm=np.zeros((ny,nx,3),dtype='uint8')

#Flatterned arrays
nxny=nx*ny
ix_f=np.reshape(ix,nxny)
iy_f=np.reshape(iy,nxny)
c_f=np.reshape(c,nxny)
z_f=c_f.copy()   #the iterated variable

for iter in range(max_iter):
    if not len(z_f):   #all points have escaped
        break
    n=iter+1
    r,g,b=n%4*64, n%8*32, n%16*16
    z_f*=z_f
    z_f+=c_f
    #points which are escaping
    escape=np.abs(z_f)>2.0
    #Set the rgb pixel value for the escaping points
    esc_parm[iy_f[escape],ix_f[escape],:]=r,g,b
    #points not escaping
    escape=~escape
    #remove batch of newly escaped points
    ix_f=ix_f[escape]
    iy_f=iy_f[escape]
    c_f=c_f[escape]
    z_f=z_f[escape]
    
print('Time taken=',time()-start_time)

from PIL import Image
picture=Image.fromarray(esc_parm)
picture.show()
#picture.save('mandelbrot.png')
picture.save('pic/fig-5-11.png')



